/*
 * AIEngine a new generation network intrusion detection system.
 *
 * Copyright (C) 2013-2023  Luis Campo Giralte
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Library General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Library General Public License for more details.
 *
 * You should have received a copy of the GNU Library General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301, USA.
 *
 * Written by Luis Campo Giralte <luis.camp0.2009@gmail.com>
 *
 */
#include "StackLan.h"

namespace aiengine {

StackLan::StackLan() {

	name("Lan network stack");

	// Add the specific Protocol object
	addProtocol(eth, true);
	addProtocol(vlan, false);
	addProtocol(mpls, false);
	addProtocol(pppoe, false);
	addProtocol(ip_, true);
	addProtocol(tcp_, true);
	addProtocol(udp_, true);
	addProtocol(icmp_, true);

	// Add the layer7 protocols in order to show pretty output
        addProtocol(http);
        addProtocol(ssl);
        addProtocol(smtp);
        addProtocol(imap);
        addProtocol(pop);
        addProtocol(ssh);
        addProtocol(bitcoin);
        addProtocol(modbus);
        addProtocol(mqtt);
        addProtocol(smb);
        addProtocol(dcerpc);
        addProtocol(tcp_generic);
        addProtocol(freqs_tcp, false);

        addProtocol(dns);
        addProtocol(sip);
        addProtocol(dhcp);
        addProtocol(ntp);
        addProtocol(snmp);
        addProtocol(ssdp);
	addProtocol(netbios);
        addProtocol(coap);
        addProtocol(rtp);
        addProtocol(quic);
        addProtocol(dtls);
        addProtocol(udp_generic);
        addProtocol(freqs_udp, false);

	// Link the FlowCaches to their corresponding FlowManager for timeouts
	flow_table_udp_->setFlowCache(flow_cache_udp_);
	flow_table_tcp_->setFlowCache(flow_cache_tcp_);

	// Connect the Protocols with the Multiplexers
	ip_->setMultiplexer(mux_ip);
	mux_ip->setProtocol(static_cast<ProtocolPtr>(ip_));

	icmp_->setMultiplexer(mux_icmp_);
	mux_icmp_->setProtocol(static_cast<ProtocolPtr>(icmp_));

	udp_->setMultiplexer(mux_udp_);
	mux_udp_->setProtocol(static_cast<ProtocolPtr>(udp_));
	ff_udp_->setProtocol(static_cast<ProtocolPtr>(udp_));

	tcp_->setMultiplexer(mux_tcp_);
	mux_tcp_->setProtocol(static_cast<ProtocolPtr>(tcp_));
	ff_tcp_->setProtocol(static_cast<ProtocolPtr>(tcp_));

	// configure the multiplexers
	mux_eth->addUpMultiplexer(mux_ip);
	mux_ip->addDownMultiplexer(mux_eth);
	mux_ip->addUpMultiplexer(mux_udp_);
	mux_udp_->addDownMultiplexer(mux_ip);
	mux_ip->addUpMultiplexer(mux_tcp_);
	mux_tcp_->addDownMultiplexer(mux_ip);
	mux_ip->addUpMultiplexer(mux_icmp_);
	mux_icmp_->addDownMultiplexer(mux_ip);

	// Connect the FlowManager and FlowCache
	tcp_->setFlowCache(flow_cache_tcp_);
	tcp_->setFlowManager(flow_table_tcp_);
	flow_table_tcp_->setProtocol(tcp_);

	udp_->setFlowCache(flow_cache_udp_);
	udp_->setFlowManager(flow_table_udp_);
	flow_table_udp_->setProtocol(udp_);

	// Connect to upper layers the FlowManager for release caches
	http->setFlowManager(flow_table_tcp_);
	ssl->setFlowManager(flow_table_tcp_);
	smtp->setFlowManager(flow_table_tcp_);
	imap->setFlowManager(flow_table_tcp_);
	pop->setFlowManager(flow_table_tcp_);
	bitcoin->setFlowManager(flow_table_tcp_);
	mqtt->setFlowManager(flow_table_tcp_);
	smb->setFlowManager(flow_table_tcp_);
	ssh->setFlowManager(flow_table_tcp_);
	dcerpc->setFlowManager(flow_table_tcp_);

	dns->setFlowManager(flow_table_udp_);
	sip->setFlowManager(flow_table_udp_);
	ssdp->setFlowManager(flow_table_udp_);
	coap->setFlowManager(flow_table_udp_);
	netbios->setFlowManager(flow_table_udp_);
	dhcp->setFlowManager(flow_table_udp_);
	quic->setFlowManager(flow_table_udp_);
	dtls->setFlowManager(flow_table_udp_);

	freqs_tcp->setFlowManager(flow_table_tcp_);
	freqs_udp->setFlowManager(flow_table_udp_);

	// Connect the AnomalyManager with the protocols that may have anomalies
	ip_->setAnomalyManager(anomaly_);
	tcp_->setAnomalyManager(anomaly_);
	udp_->setAnomalyManager(anomaly_);

	// Configure the FlowForwarders
	tcp_->setFlowForwarder(ff_tcp_);
	udp_->setFlowForwarder(ff_udp_);

	// List of default protocols enabled on this stack
	setTCPDefaultForwarder(ff_tcp_);
	setUDPDefaultForwarder(ff_udp_);

	std::ostringstream msg;

        msg << name() << " ready.";

        infoMessage(msg.str());

	setMode("full");
}

void StackLan::showFlows(std::basic_ostream<char> &out, std::function<bool (const Flow&)> condition, int protocol) const {

        int total = flow_table_tcp_->getTotalFlows() + flow_table_udp_->getTotalFlows();
        out << "Flows on memory " << total << std::endl;
	if (protocol == IPPROTO_TCP)
		flow_table_tcp_->showFlows(out, condition);
	else if (protocol == IPPROTO_UDP)
		flow_table_udp_->showFlows(out, condition);
	else {
		flow_table_tcp_->showFlows(out, condition);
		flow_table_udp_->showFlows(out, condition);
	}
}

void StackLan::showFlows(Json &out, std::function<bool (const Flow&)> condition, int protocol) const {

	if (protocol == IPPROTO_TCP)
		flow_table_tcp_->showFlows(out, condition);
	else if (protocol == IPPROTO_UDP)
		flow_table_udp_->showFlows(out, condition);
	else {
		flow_table_tcp_->showFlows(out, condition);
		flow_table_udp_->showFlows(out, condition);
	}
}

void StackLan::setMode(const std::string &mode) {

	std::initializer_list<SharedPointer<FlowForwarder>> tcp_list {
		ff_tcp_, ff_http, ff_ssl, ff_smtp, ff_imap, ff_pop, ff_bitcoin, ff_modbus,
		ff_mqtt, ff_smb, ff_ssh, ff_dcerpc, ff_tcp_generic, ff_tcp_freqs};
	std::initializer_list<SharedPointer<FlowForwarder>> udp_list {
		ff_udp_, ff_dns, ff_sip, ff_dhcp, ff_ntp, ff_snmp, ff_ssdp, ff_netbios,
		ff_coap, ff_rtp, ff_quic, ff_dtls, ff_udp_generic, ff_udp_freqs};

	if (auto it = modes.find(mode); it != modes.end()) {
		std::ostringstream msg;

		disableFlowForwarders(tcp_list);
		disableFlowForwarders(udp_list);

		if ((*it).second == Mode::FULL) {

			operation_mode_ = "full";
			msg << "Enable FullEngine mode on " << name();
			enableFlowForwarders(tcp_list);
			enableFlowForwarders(udp_list);

		} else if ((*it).second == Mode::FREQUENCIES) {

			operation_mode_ = "frequency";
			msg << "Enable FrequencyEngine mode on " << name();
			enableFlowForwarders({ff_tcp_, ff_tcp_freqs});
			enableFlowForwarders({ff_udp_, ff_udp_freqs});

		} else if ((*it).second == Mode::NIDS) {

			operation_mode_ = "nids";
			msg << "Enable NIDSEngine mode on " << name();
			enableFlowForwarders({ff_tcp_, ff_tcp_generic});
			enableFlowForwarders({ff_udp_, ff_udp_generic});
		}
		infoMessage(msg.str());
	}
}

void StackLan::setTotalTCPFlows(int value) {

	flow_cache_tcp_->create(value);
	tcp_->createTCPInfos(value);

	// The vast majority of the traffic of internet is HTTP
	// so create 75% of the value received for the http caches
	http->increaseAllocatedMemory(value * 0.75);

	// The 40% of the traffic is SSL
	ssl->increaseAllocatedMemory(value * 0.4);

        // 5% of the traffic could be SMTP/IMAP, im really positive :D
        smtp->increaseAllocatedMemory(value * 0.05);
        imap->increaseAllocatedMemory(value * 0.05);
        pop->increaseAllocatedMemory(value * 0.05);
        bitcoin->increaseAllocatedMemory(value * 0.05);
        mqtt->increaseAllocatedMemory(value * 0.05);
        smb->increaseAllocatedMemory(value * 0.05);
        ssh->increaseAllocatedMemory(value * 0.05);
        dcerpc->increaseAllocatedMemory(value * 0.05);
        freqs_tcp->increaseAllocatedMemory(16);
}

void StackLan::setTotalUDPFlows(int value) {

	flow_cache_udp_->create(value);

	dns->increaseAllocatedMemory(value / 2);
	sip->increaseAllocatedMemory(value * 0.2);
	ssdp->increaseAllocatedMemory(value * 0.2);
	coap->increaseAllocatedMemory(value * 0.2);
	netbios->increaseAllocatedMemory(value * 0.2);
	dhcp->increaseAllocatedMemory(value * 0.1);
	quic->increaseAllocatedMemory(value * 0.2);
	dtls->increaseAllocatedMemory(value * 0.2);
        freqs_udp->increaseAllocatedMemory(16);
}

int StackLan::getTotalTCPFlows() const { return flow_cache_tcp_->getTotalFlows(); }

int StackLan::getTotalUDPFlows() const { return flow_cache_udp_->getTotalFlows(); }

void StackLan::setFlowsTimeout(int timeout) {

        flow_table_udp_->setTimeout(timeout);
        flow_table_tcp_->setTimeout(timeout);
}

void StackLan::setTCPRegexManager(const SharedPointer<RegexManager> &rm) {

	tcp_->setRegexManager(rm);
	tcp_generic->setRegexManager(rm);
	super_::setTCPRegexManager(rm);
}

void StackLan::setUDPRegexManager(const SharedPointer<RegexManager> &rm) {

	udp_->setRegexManager(rm);
	udp_generic->setRegexManager(rm);
	super_::setUDPRegexManager(rm);
}

void StackLan::setTCPIPSetManager(const SharedPointer<IPSetManager> &ipset_mng) {

	tcp_->setIPSetManager(ipset_mng);
	super_::setTCPIPSetManager(ipset_mng);
}

void StackLan::setUDPIPSetManager(const SharedPointer<IPSetManager> &ipset_mng) {

	udp_->setIPSetManager(ipset_mng);
	super_::setUDPIPSetManager(ipset_mng);
}

void StackLan::stopAsioService() {

	// The sockets created by the RejectManager
	// needs to be reset them in order to choose other io_service
#if defined(HAVE_REJECT_FLOW)
	rj_mng_->close();
#endif
}

void StackLan::setAsioService(boost::asio::io_service &io_service, const std::string &device) {

	// Create a new RejectManager with their corresponding sockets
#if defined(HAVE_REJECT_FLOW)
	if (geteuid() == 0) { // The process have rights on raw sockets
		rj_mng_->setAsioService(io_service, device);
		if (rj_mng_->ready()) {
			// Attach the reject function to the corresponding protocols tcp/udp
			tcp_->addRejectFunction(std::bind(&RejectManager<StackLan>::rejectTCPFlow, rj_mng_, std::placeholders::_1));
			udp_->addRejectFunction(std::bind(&RejectManager<StackLan>::rejectUDPFlow, rj_mng_, std::placeholders::_1));
		}
	}
#endif
}

void StackLan::statistics(std::basic_ostream<char> &out) const {

	super_::statistics(out);

#ifdef HAVE_REJECT_FLOW
	if (geteuid() == 0) { // The process have rights on raw sockets
		if (rj_mng_) {
			rj_mng_->statistics(out);
			out << std::endl;
		}
	}
#endif
}

#if defined(JAVA_BINDING)

void StackLan::setTCPRegexManager(RegexManager *sig) {

	SharedPointer<RegexManager> rm;

	if (sig != nullptr)
		rm.reset(sig);
	else
		rm.reset();

	setTCPRegexManager(rm);
}

void StackLan::setUDPRegexManager(RegexManager *sig) {

	SharedPointer<RegexManager> rm;

	if (sig != nullptr)
		rm.reset(sig);
	else
		rm.reset();

	setUDPRegexManager(rm);
}

void StackLan::setTCPIPSetManager(IPSetManager *ipset_mng) {

	SharedPointer<IPSetManager> im;

	if (ipset_mng != nullptr)
		im.reset(ipset_mng);
	else
		im.reset();

	setTCPIPSetManager(im);
}

void StackLan::setUDPIPSetManager(IPSetManager *ipset_mng) {

	SharedPointer<IPSetManager> im;

	if (ipset_mng != nullptr)
		im.reset(ipset_mng);
	else
		im.reset();

	setUDPIPSetManager(im);
}

#endif

std::tuple<Flow*, Flow*> StackLan::getCurrentFlows() const {

	Flow *flow = nullptr;
	uint16_t proto = ip_->getProtocol();

	if (proto == IPPROTO_TCP)
		flow = tcp_->getCurrentFlow();
	else if (proto == IPPROTO_UDP)
		flow = udp_->getCurrentFlow();

#if GCC_VERSION < 50500
	return std::tuple<Flow*, Flow*>(flow, nullptr);
#else
	return {flow, nullptr};
#endif
}

SharedPointer<Flow> StackLan::getFlow(const FlowSearchOptions &fsos) const {

	if (fsos.protocol == IPPROTO_TCP)
		return flow_table_tcp_->find(fsos);
	else
		return flow_table_udp_->find(fsos);
}

} // namespace aiengine
